#pragma region Copyright
/// <copyright>
/// Copyright (c) 2007 Ilia Shramko.
/// Released under the terms of the GNU Lesser General Public License.
/// </copyright>
#pragma endregion

#pragma once

#include <atlbase.h>
#include <atlwin.h>
#include <vector>


// {4498612A-6EA9-4983-88F4-1385EBAFB7AD}
DEFINE_GUID(CLSID_SimpleGrabber,
0x4498612a, 0x6ea9, 0x4983, 0x88, 0xf4, 0x13, 0x85, 0xeb, 0xaf, 0xb7, 0xad);

template<class TClass, int MessageMapId = 0>
class WindowMessageSink : public CWindowImpl<WindowMessageSink<TClass>, CWindow, CFrameWinTraits >
{
private:
	TClass* _client;

	BEGIN_MSG_MAP(WindowMessageSink)
		CHAIN_MSG_MAP_ALT_MEMBER((*_client), MessageMapId)
	END_MSG_MAP()

public:

	WindowMessageSink(TClass* client)
	{
		_client = client;
	}

};


class CSimpleGrabberInPin;
class CSimpleGrabberOutPin;
class CSimpleGrabber;
class CSimpleGrabberMemAllocator;


//----------------------------------------------------------------------------
//  CSimpleGrabberMemAllocator
//----------------------------------------------------------------------------
class CSimpleGrabberMemAllocator : public CMemAllocator
{
private:
	CMediaType _mediaType;
	bool _mediaTypeChanged;
	bool _upstreamFilterNotified;

public:
	CSimpleGrabberMemAllocator(HRESULT* hr) :
		CMemAllocator(NAME("CSGMemAllocator"), NULL, hr), _mediaTypeChanged(false), _upstreamFilterNotified(false)
	{
	}

    STDMETHOD(GetBuffer)(IMediaSample** ppBuffer,
		REFERENCE_TIME* pStartTime, REFERENCE_TIME* pEndTime, DWORD dwFlags);

	HRESULT GetMediaType(CMediaType& mediaType) const;
	HRESULT SetMediaType(const AM_MEDIA_TYPE& mt);
};


//----------------------------------------------------------------------------
//  CSimpleGrabber
//----------------------------------------------------------------------------
class CSimpleGrabber : public CTransformFilter
{
private:
	enum {
		CommandMessageSinkMapID = 1,
		WM_APP_SAVESCREENSHOT = WM_APP + 1,
		WM_APP_SHUTDOWN = WM_APP + 2,
		STATE_WAITING = 0,
		STATE_PREPARING = 1,
		STATE_SAVING = 2,
	};

	WindowMessageSink<CSimpleGrabber, CommandMessageSinkMapID> _commandMessageSink;
	int _takeScreenshotState;
	DWORD _takeScreenshotRequestTime;
	std::wstring _screenshotFilePath;
	CMediaType _sampleMediaType;
	BYTE* _sampleBuffer;
	size_t _sampleBufferSize;

	CSimpleGrabberInPin* GetInput();
	CSimpleGrabberOutPin* GetOutput();

	LRESULT OnCommandWindowSetText(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& bHandled);
	LRESULT OnSaveScreenshot(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& bHandled);
	LRESULT OnSaveShutDown(UINT uMsg, WPARAM wParam, LPARAM lParam, BOOL& bHandled);
	
	HRESULT PrepareScreenshot(IMediaSample* mediaSample);
	HRESULT SaveScreenshot();
	HRESULT SaveBitmapFile(LPCWSTR filePath,
		const BYTE* bitmapBits, size_t bitmapBitsSize, int width, int height, WORD colorMode);
	HRESULT SaveBitmapFile(LPCWSTR filePath, const BITMAPINFO* bitmapInfo);

	HRESULT Copy(IMediaSample* pSource, IMediaSample* pDest) const;
	HRESULT CopyBuffer(const AM_MEDIA_TYPE* srcMediaType, const BYTE* srcBuffer, size_t srcBufferSize,
		AM_MEDIA_TYPE* dstMediaType, BYTE* dstBuffer, size_t dstBufferSize) const;

protected:

    CMediaType m_mtAccept;
    CCritSec m_Lock; // serialize access to our data

public:

    CSimpleGrabber(IUnknown* pOuter, HRESULT* pHr);
	~CSimpleGrabber();

    static CUnknown* WINAPI CreateInstance(LPUNKNOWN punk, HRESULT *phr);

    // Expose ISimpleGrabber
    STDMETHODIMP NonDelegatingQueryInterface(REFIID riid, void ** ppv);
    DECLARE_IUNKNOWN;

	BEGIN_MSG_MAP(CSimpleGrabber)
		ALT_MSG_MAP(CommandMessageSinkMapID)
			MESSAGE_HANDLER(WM_SETTEXT, OnCommandWindowSetText)
			MESSAGE_HANDLER(WM_APP_SAVESCREENSHOT, OnSaveScreenshot)
			MESSAGE_HANDLER(WM_APP_SHUTDOWN, OnSaveShutDown)
	END_MSG_MAP()

	STDMETHOD(JoinFilterGraph)(IFilterGraph* pGraph, LPCWSTR pName);

    virtual HRESULT CheckInputType(const CMediaType * pmt);
    virtual HRESULT CheckTransform(const CMediaType* mtIn, const CMediaType* mtOut);
    virtual HRESULT DecideBufferSize(IMemAllocator* pAllocator, ALLOCATOR_PROPERTIES* pprop);
    virtual HRESULT GetMediaType(int iPosition, CMediaType* pMediaType);
    virtual HRESULT CompleteConnect(PIN_DIRECTION direction, IPin* pReceivePin);
    virtual HRESULT Transform(IMediaSample* pIn, IMediaSample* pOut);
	
	HRESULT QueryAcceptOutput(const AM_MEDIA_TYPE* pmt);

	static HRESULT AreMediaTypesCompatible(const AM_MEDIA_TYPE* mtIn, const AM_MEDIA_TYPE* mtOut);

};


//----------------------------------------------------------------------------
// we override the input pin class so we can provide a media type
// to speed up connection times. When you try to connect a filesourceasync
// to a transform filter, DirectShow will insert a splitter and then
// start trying codecs, both audio and video, video codecs first. If
// your sample grabber's set to connect to audio, unless we do this, it
// will try all the video codecs first. Connection times are sped up x10
// for audio with just this minor modification!
//----------------------------------------------------------------------------
class CSimpleGrabberInPin : public CTransformInputPin
{
private:
	std::vector<AM_MEDIA_TYPE*> _upFilterMediaTypes;

	HRESULT LoadMediaTypes(IPin* pin);
	void ClearMediaTypes();

public:

	CSimpleGrabberInPin(CSimpleGrabber* pFilter, HRESULT* pHr)
		: CTransformInputPin(TEXT("SimpleGrabberInputPin\0"), pFilter, pHr, L"Input\0")
	{
	}

    virtual ~CSimpleGrabberInPin();

	STDMETHOD(GetAllocator)(IMemAllocator** ppAllocator);
    STDMETHOD(NotifyAllocator)(IMemAllocator* pAllocator, BOOL bReadOnly);

	STDMETHOD(ReceiveConnection)(IPin* pConnector, const AM_MEDIA_TYPE* pmt);

	// override to provide major media type for fast connects
    virtual HRESULT GetMediaType(int iPosition, CMediaType* pMediaType);

	CSimpleGrabberMemAllocator* GetAllocator() { return static_cast<CSimpleGrabberMemAllocator*>(m_pAllocator); }

	std::vector<AM_MEDIA_TYPE*>& GetUpFilterMediaTypes() { return _upFilterMediaTypes; }
	HRESULT CanReconnect(const CMediaType& outputMediaType, CMediaType* upFilterMediaType);
    HRESULT GetUpFilterMediaType(int iPosition, CMediaType* pMediaType);
};

class CSimpleGrabberOutPin : public CTransformOutputPin
{
public:
    CSimpleGrabberOutPin(CSimpleGrabber* pFilter, HRESULT* pHr)
		: CTransformOutputPin(TEXT("SimpleGrabberOutputPin\0"), pFilter, pHr, L"Output\0")
	{
	}

	STDMETHOD(QueryAccept)(const AM_MEDIA_TYPE* pmt);
};


